package com.alibaba.datax.plugin.writer.spdmysqlwriter;

import com.alibaba.datax.common.base.IdRecord;
import com.alibaba.datax.common.constant.IdRecordConstants;
import com.alibaba.datax.common.element.Record;
import com.alibaba.datax.common.plugin.RecordReceiver;
import com.alibaba.datax.common.spi.Writer;
import com.alibaba.datax.common.util.Configuration;
import com.alibaba.datax.common.util.RedissonUtils;
import com.alibaba.datax.plugin.rdbms.util.DBUtil;
import com.alibaba.datax.plugin.rdbms.util.DataBaseType;
import com.alibaba.datax.plugin.rdbms.writer.Key;
import com.alibaba.fastjson2.JSONObject;
import org.redisson.api.RTopic;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;


//TODO writeProxy
public class SpdMysqlWriter extends Writer {
    private static final DataBaseType DATABASE_TYPE = DataBaseType.MySql;

    public static class Job extends Writer.Job {
        private Configuration originalConfig = null;

        @Override
        public void preCheck(){
            this.init();
        }

        @Override
        public void init() {
            this.originalConfig = super.getPluginJobConf();
        }

        // 一般来说，是需要推迟到 task 中进行pre 的执行（单表情况例外）
        @Override
        public void prepare() {
        }

        @Override
        public List<Configuration> split(int mandatoryNumber) {
            List<Configuration> configurations = new ArrayList<>();
            configurations.add(this.originalConfig);
            return configurations;
        }

        // 一般来说，是需要推迟到 task 中进行post 的执行（单表情况例外）
        @Override
        public void post() {
        }

        @Override
        public void destroy() {
        }

    }

    public static class Task extends Writer.Task {
        private static final Logger LOG = LoggerFactory
                .getLogger(Task.class);
        private Configuration writerSliceConfig;
        private Connection connection;

        private String insertSql;
        @Override
        public void init() {
            this.writerSliceConfig = super.getPluginJobConf();
            String username = writerSliceConfig.getString(Key.USERNAME);
            String password = writerSliceConfig.getString(Key.PASSWORD);
            String jdbcUrl = writerSliceConfig.getString(Key.JDBC_URL);
            insertSql = writerSliceConfig.getString("insertSql");
            connection = DBUtil.getConnection(DATABASE_TYPE,jdbcUrl, username, password);
        }

        @Override
        public void prepare() {
        }

        //TODO 改用连接池，确保每次获取的连接都是可用的（注意：连接可能需要每次都初始化其 session）
        public void startWrite(RecordReceiver recordReceiver) {
            String redisHost = System.getProperty("redisHost");
            String redisPasswd = System.getProperty("redisPasswd");
            String redisDatabase = System.getProperty("redisDatabase");
            String taskId = System.getProperty("taskId");
            String incrementField = System.getProperty("incrementField");

            LOG.info("开始写入数据任务:{}",taskId);
            RedissonUtils redissonUtils = new RedissonUtils(redisHost, redisPasswd, 6379, Integer.parseInt(redisDatabase));
            RTopic idRecord = redissonUtils.getRedissonClient().getTopic("idRecord");

            Record record;
            try {
                int count = 0;
                connection.setAutoCommit(false);
                PreparedStatement preparedStatement = connection.prepareStatement(insertSql);
                String lastVal = "";
                while ((record = recordReceiver.getFromReader()) != null) {
                    Map<String, String> meta = record.getMeta();
                    lastVal = meta.get(incrementField);
                    count++;
                    for (int i=0;i<record.getColumnNumber();i++){
                        preparedStatement.setString(i+1,record.getColumn(i).asString());
                    }
                    preparedStatement.addBatch();
                    if(count>=1000){
                        preparedStatement.executeBatch();
                        preparedStatement.clearBatch();
                        count = 0;
                        connection.commit();//执行完后，手动提交事务
                        idRecord.publish(JSONObject.toJSONString(new IdRecord(taskId,lastVal, IdRecordConstants.WRITE)));
                    }
                }
                preparedStatement.executeBatch();
                connection.commit();
                connection.setAutoCommit(true);
                idRecord.publish(JSONObject.toJSONString(new IdRecord(taskId,lastVal, IdRecordConstants.WRITE)));
            }catch (Exception e){
               e.printStackTrace();
            }finally {
                DBUtil.closeDBResources(null, null, connection);
            }

        }

        @Override
        public void post() {
        }

        @Override
        public void destroy() {
        }

        @Override
        public boolean supportFailOver(){
            String writeMode = writerSliceConfig.getString(Key.WRITE_MODE);
            return "replace".equalsIgnoreCase(writeMode);
        }

    }


}
